{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MEng Project"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Oct 24th Meeting Action Items:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- [x] Always show the source signal (check for artifacts) and keep length of time for signal consistent (~3 min)\n",
    "- [x] Double check fantasia source collection\n",
    "   - Breathing signal collected using respiratory belt\n",
    "- [x] FIX PEAK DETECTION AND INTERP (Look at paper)\n",
    "   - martinez2004 is fastest and most reliable peak detection\n",
    "- [x] CWT on Respiration for same 3 patients, CWT on ECG of patients\n",
    "- [x] Stacked figure, resp (for signal we just looked at), ecg_env, ecg, CWT (indicate which points were used for interpolation)\n",
    "- [x] Older patient\n",
    "- [ ] Ask GPT about some ideas\n",
    "- [ ] wavelet phase coherence between envelope and the IP signal. Because we don't care about amplitude we care about frequencies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Nov 7th Meeting Action Items:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- [x] Wavelet Phase Coherence revisit code and plot with gaussian blur\n",
    "- [ ] Phase Amplitude Coupling on ECG (hi freq) (should find what is the low frequency that modulates the high frequency)\n",
    "- [x] Ensure window size is appropriate for each frequency and Ensure frequency range for graph is appropriate (low should cover respiratory range (varies with dataset), high should span ECG range (~1hz for R-R, Q-T is a bit slower. Stay with R-R for now.))\n",
    "- [x] ML for respiration signal from ECG"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Meeting Nov 21st\n",
    "- [ ] split data into train and test\n",
    "- [ ] choose MARSH sample for testing. see the difference between vanilla xgb and vanilla + marsh trained data xgb\n",
    "- [ ] add pink noise (establish tolerance) look at chris's paper\n",
    "- [ ] wavelet phase coherence between extracted and ground truth (marija cotic paper, [here](https://onlinelibrary.wiley.com/doi/full/10.1111/epi.12918))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Meeting Jan 9th, 2024\n",
    "- [x] IEEE Format, check website. 46th anual IEEE EMBC 2024\n",
    "- [x] Double check that Respiration signal and \"rate\" are not conflated\n",
    "- [x] Get PAC's PCC and coherence on Marsh Data (Adam)\n",
    "- [ ] Add PCC metrics to figure 10  (Figure 3 in new paper)\n",
    "- [x] Show samples of data, combined MARSH (figure 1)\n",
    "- [x] Figure 7, 1 minute, with all 3 methods (Augusto) (Figure 2 in new Paper)\n",
    "- [x] Summarize methods, save space for discussions.\n",
    "- [x] WPC for each candidate A ONLY (no IMF stuff) against the ground truth.\n",
    "- [ ] t-test between methods\n",
    "- [x] Normalized graph (include in text)\n",
    "- [x] Use annotated peaks (if time permits)\n",
    "- [x] Remove IMF stuffs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Meeting Jan 22nd, 2024\n",
    "- [x] Figures from the same subject, same time segment (subject 0023, minute 1:00-2:00)\n",
    "- [x] same # of segments, same # of patients. Align completely\n",
    "- [x] send adam ALL the ECG envelope data.\n",
    "- [x] Take 2 chirp signals + AWGN, CWT, WPC them and validate if PC_norm or PC is better.\n",
    "- [x] Deep dive into a single sample (7113) to debug our approach for one minute, and two minute combinations\n",
    "- [x] See if there is a difference between results for all minute 1's, minute 2's, minute 3's.\n",
    "- [x] Once we have the results, find a sample for our comparison plot that is representative of the results.\n",
    "- [x] Remove IMF from figure 3 in paper.\n",
    "- [x] Take out bottom subfigure for ECG/IP example to highlight the PAC motivation\n",
    "\n",
    "# Meeting Jan 29th, 2024\n",
    "- [ ] Take out all mentions of PCC\n",
    "- [x] Error bars for WPC graph\n",
    "- [x] Get plot 6 color consistency (ECG Blue, PAC Orange) and use those colors in the bar graph\n",
    "- [ ] Discussion points"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Meeting Feb 26th\n",
    "- [x] Set up infrastructure for ML experiment tracking\n",
    "- [x] Retrain Models on first 3 minutes\n",
    "\n",
    "# Meeting March 4th AND 11th\n",
    "- [ ] Reliability of IP versus Respiratory Belt\n",
    "- [ ] Get score of WPC of 10s, 20s\n",
    "- [ ] Deep Learning\n",
    "- [ ] See if leaving out bad examples helps\n",
    "- [x] Analyze from 420s to 480s with 0.133hz\n",
    "- [x] Why is the WPC greater than 1 sometimes? Reimplement?\n",
    "- [ ] Estimation statistics, comparing instantaneous frequencies between prediction/ground truth to also estimate the goodness of results next to WPC. We can't use straight lines for these nonlinear systems"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Meeting April 15th\n",
    "- [x] Fig. 1: Reduce line thickness, align axes. Aim for 300 DPI, as svg\n",
    "- [x] Find a ground truth signal for figure 1 where the amplitudes are more uniform (subject 0047, minute 1. crop to valid points)\n",
    "- [ ] replace fig 1 w/ corresponing subject (all 300 DPI)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adam PAC Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install scipy neurokit2 numpy plotly matplotlib tqdm EMD-signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.io import loadmat\n",
    "# X7113 = loadmat('../MARSH/subj7113_cwt_respiration.mat')\n",
    "# plt.plot(X7113['pt7113'].T)\n",
    "\n",
    "loaded_data = loadmat('../MARSH/extracted.mat')\n",
    "PAC_signals = loaded_data['extracted_signals']\n",
    "PAC_subjects = loaded_data[\"subject_numbers\"]\n",
    "PAC_ids = loaded_data[\"extracted_signal_ids\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scipy, EMD, and numpy Transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy\n",
    "from scipy.signal import butter, filtfilt, detrend\n",
    "from PyEMD import EMD, EEMD\n",
    "import neurokit2 as nk\n",
    "import numpy as np\n",
    "import plotly.graph_objects as go\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "def FFT(x, fs=250, cutoff=np.inf):\n",
    "    X = np.fft.fft(x)\n",
    "    F = np.fft.fftfreq(x.shape[0], d=1/fs)\n",
    "    return F[(F>0) & (F<cutoff)], abs(X[(F>0) & (F<cutoff)])\n",
    "\n",
    "# Function to design a Butterworth bandpass filter\n",
    "def butter_bandpass(lowcut, highcut, fs, order=5):\n",
    "    nyq = 0.5 * fs\n",
    "    low = lowcut / nyq\n",
    "    high = highcut / nyq\n",
    "    b, a = butter(order, [low, high], btype='band')\n",
    "    return b, a\n",
    "\n",
    "# Function to apply the bandpass filter\n",
    "def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):\n",
    "    b, a = butter_bandpass(lowcut, highcut, fs, order=order)\n",
    "    y = filtfilt(b, a, data)\n",
    "    return y\n",
    "\n",
    "# Function to design a Butterworth lowpass filter\n",
    "def butter_lowpass(cutoff, fs, order=5):\n",
    "    nyq = 0.5 * fs\n",
    "    normal_cutoff = cutoff / nyq\n",
    "    b, a = butter(order, normal_cutoff, btype='low', analog=False)\n",
    "    return b, a\n",
    "\n",
    "# Function to apply the lowpass filter\n",
    "def butter_lowpass_filter(data, cutoff, fs, order=5):\n",
    "    b, a = butter_lowpass(cutoff, fs, order=order)\n",
    "    y = filtfilt(b, a, data)\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_signal_extrema(signal, maxima=True):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    \n",
    "    # 1st derivative, if close to zero, then the point is an extrema\n",
    "    dxdt = np.diff(signal, prepend=signal[0])\n",
    "    atol = (signal.max()-signal.min())*1e-3\n",
    "    # print(\"Min Diff and tolerance\", abs(dxdt[1:]).min(), atol)\n",
    "\n",
    "    # 2nd der, if negative, point is concave down, else concave up\n",
    "    dx2dt2 = np.diff(dxdt, prepend=dxdt[0])\n",
    "\n",
    "    # check if extrema, concave down, and positive signal\n",
    "    if maxima:\n",
    "        peak_bools = np.isclose(dxdt, 0, atol=atol) & (dx2dt2 < 0) & (signal > 0)\n",
    "    # check if extrema, concave up, and negative signal\n",
    "    else:\n",
    "        peak_bools = np.isclose(dxdt, 0, atol=atol) & (dx2dt2 > 0) & (signal < 0)\n",
    "        \n",
    "    peak_idxs = np.arange(signal.shape[0])[peak_bools]\n",
    "    \n",
    "    # ensure distance between peak idxs at least 2x0.55 hz worth of samples\n",
    "    filter = np.diff(peak_idxs, prepend=peak_idxs[0]-1) >= 250/0.55 # REMOVED (2*0.55)\n",
    "    peak_idxs = peak_idxs[filter]\n",
    "    \n",
    "    peaks = signal[peak_idxs]\n",
    "    return peak_idxs, peaks\n",
    "\n",
    "\n",
    "def zero_crossings(signal):\n",
    "    raise NotImplementedException\n",
    "    return\n",
    "    # return np.where(np.diff(np.sign(signal)))[0]\n",
    "\n",
    "def ins_freqs(signal, ref_pts, fs=250): # REMOVED FIRST ARG\n",
    "# def ins_freqs(ref_pts, fs=250):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "    - signal: input signal\n",
    "    - ref_pts: refrence_points. the intervals between these points will be used to calculate the instantaneous frequencies\n",
    "    \"\"\"\n",
    "    wavelengths = np.diff(ref_pts)\n",
    "    return ref_pts/fs, fs/wavelengths\n",
    "\n",
    "def eval(t_gt, if_gt, t_est, if_est, plot=False):\n",
    "    \"\"\"\n",
    "    t_gt, if_gt: time and instantaneous frequencies from ground truth\n",
    "    t_est, if_est: time and instantaneous frequencies from estimated breathing signal\n",
    "    \"\"\"\n",
    "    \n",
    "    alignment = np.abs(\n",
    "        (t_gt[1:, np.newaxis]-t_est[1:])\n",
    "    ).argmin(axis=1)\n",
    "    if plot:\n",
    "        plt.plot(\n",
    "            [0, 1],\n",
    "            [0, 1],\n",
    "            '-'\n",
    "        )\n",
    "        plt.plot(\n",
    "            if_gt,\n",
    "            if_est[alignment],\n",
    "            'bo'\n",
    "        )\n",
    "        plt.xlabel('Ground Truth Instantaneous Frequencies')\n",
    "        plt.ylabel('Estimate Instantaneous Frequencies')\n",
    "        plt.legend()\n",
    "        plt.show()\n",
    "\n",
    "    return if_gt.mean(), if_est[alignment].mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from Data import MarshData, FantasiaData\n",
    "from Transforms import WPC\n",
    "import Transforms\n",
    "from Signal import Signal\n",
    "\n",
    "marsh_dataset = [\n",
    "    MarshData(f\"../MARSH/{i}/\") for i in os.listdir(\"../MARSH/\") if len(i) == 4\n",
    "]\n",
    "\n",
    "exclusion_list = [\n",
    "    ('1436', 2),\n",
    "    ('0046', 0),\n",
    "    ('2655', 0),\n",
    "    ('2655', 2),\n",
    "    ('0048', 2),\n",
    "    ('0037', 1),\n",
    "    ('5329', 0)\n",
    "]\n",
    "\n",
    "m1 = MarshData(\"../MARSH/0023/\")\n",
    "m2 = MarshData(\"../MARSH/7113/\")\n",
    "m_0047 = MarshData(\"../MARSH/0047/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Marsh Data Plot (fig 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "marsh_transforms = [Transforms.Crop(start=60, end=120), Transforms.MinMaxScale()]\n",
    "# fig = m1.ECG().transform(transforms=marsh_transforms).plot(transformed=True)\n",
    "# fig = m1.IP().transform(transforms=marsh_transforms).plot(base_fig=fig, transformed=True)\n",
    "# fig = m1.ECG_ENV().transform(transforms=marsh_transforms).plot(base_fig=fig, transformed=True)\n",
    "# fig.update_layout(plot_bgcolor='white')\n",
    "# fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import matplotlib.pyplot as plt\n",
    "# import numpy as np\n",
    "# # fig, ax = plt.subplots(1, 1, figsize=(10, 3.5))\n",
    "x = m_0047.ECG().transform(transforms=marsh_transforms).transformed_data\n",
    "# x -= x.mean()\n",
    "F = np.fft.fft(x)\n",
    "freq = np.fft.fftfreq(x.shape[0], d=1/250)\n",
    "# # ax.plot(60 + np.arange(0, x.shape[0]/250, 1/250), x) # because crop starts  60 seconds into recording\n",
    "# # ax.set_xlabel(\"time (s)\")\n",
    "# # ax.set_ylabel(\"Normalized Amplitude\")\n",
    "# # ax.set_title(\"MARSH: Sample ECG\")\n",
    "\n",
    "fig1 = go.Figure(\n",
    "    data = [go.Scatter(\n",
    "        x = 60 + np.arange(0, x.shape[0]/250, 1/250), \n",
    "        y = x,\n",
    "        marker=dict(color=\"rgb(124,218,252)\")\n",
    "    )]\n",
    ")\n",
    "fig1.update_layout(\n",
    "    xaxis_title=\"Time (s)\",\n",
    "    yaxis_title=\"Normalized Amplitude\",\n",
    "    # title=\"MARSH: Sample ECG\"\n",
    ")\n",
    "fig1.update_layout(plot_bgcolor='white', height=530, width=680, showlegend=False, font_family='Arial', margin=dict(l=70, r=70, t=70, b=70))\n",
    "fig1.update_xaxes(showline=True, linecolor='black', linewidth=1)\n",
    "fig1.update_yaxes(showline=True, linecolor='black', linewidth=1)\n",
    "fig1.show()\n",
    "fig1.write_image(\"paper_res/fig1.svg\")\n",
    "\n",
    "# x = m1.IP().transform(transforms=marsh_transforms).transformed_data\n",
    "# x -= x.mean()\n",
    "# F = np.fft.fft(x)\n",
    "# freq = np.fft.fftfreq(x.shape[0], d=1/250)\n",
    "# ax[1].plot(np.arange(0, x.shape[0]/250, 1/250), np.convolve(np.ones(250)/250, x, mode='same'))\n",
    "# ax[1].set_xlabel(\"time (s)\")\n",
    "# ax[1].set_title(\"MARSH: Smoothed Sample IP\")\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transforms Definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_MINUTES = 3\n",
    "pipeline = [\n",
    "        lambda x: x-x.mean(),\n",
    "        lambda x: detrend(x),\n",
    "        lambda x: np.convolve(np.ones(100)/100, x, mode='same') # exclude for method A\n",
    "]\n",
    "lo, hi, res = 0.1, 0.55, 60\n",
    "CWT = Transforms.CWT(plot=True, lower_bound=lo, higher_bound=hi, resolution=res)\n",
    "\n",
    "# different crops per minute\n",
    "# no smoothing\n",
    "method_A_transforms = [\n",
    "    [Transforms.Crop(start=60*i, end=60*(i+1))] + pipeline for i in range(NUM_MINUTES)\n",
    "]\n",
    "\n",
    "gt_transforms = [\n",
    "    [Transforms.Crop(start=60*i, end=60*(i+1))] + pipeline for i in range(NUM_MINUTES)\n",
    "]\n",
    "\n",
    "# define bands over which to calculate coherence\n",
    "freq_band = np.linspace(lo, hi, res)\n",
    "\n",
    "SAMPLE_RATE_CWT = 250"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "|Noise $\\sigma$ | PC Norm | PC |\n",
    "|  --- |  --- |  --- |\n",
    "| 0 | 1.0000 | 1.0000 |\n",
    "| 0.0125 | 0.9836 | 0.9859 |\n",
    "| 0.025 | 0.9699 | 0.9748 |\n",
    "| 0.05 | 0.9429 | 0.9531 |\n",
    "| 0.1 | 0.8978 | 0.9114 |\n",
    "| 0.5 | 0.7141 | 0.7361 |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for sigma in [0, 0.0125, 0.025, 0.05, 0.1, 0.5]:\n",
    "    t = np.arange(0, 60, 1/250)\n",
    "    f = np.linspace(lo, hi, t.shape[0])\n",
    "    clean_chirp = scipy.signal.chirp(t=t, f0=lo, t1=60, f1=hi)\n",
    "    chirpA = clean_chirp + np.random.normal(scale=sigma, size=f.shape[0])\n",
    "    chirpB = clean_chirp + np.random.normal(scale=sigma, size=f.shape[0])\n",
    "    chirp_signalA = Signal(format='mat', filepath=None, _type='ECG', data=chirpA)\n",
    "    chirp_signalB = Signal(format='mat', filepath=None, _type='ECG', data=chirpB)\n",
    "    \n",
    "    # plot\n",
    "    chirp_signalA.plot().show()\n",
    "    chirp_signalB.plot().show()\n",
    "    \n",
    "    t_test, PC_norms_test, PC_test = WPC(\n",
    "        chirp_signalA.transform(transforms=[CWT]).transformed_data,\n",
    "        chirp_signalB.transform(transforms=[CWT]).transformed_data,\n",
    "        freq=freq_band\n",
    "    )\n",
    "    print(f\"| {sigma} | {0.0:.4f} | {PC_test.mean():.4f} |\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f_ex = 0.52\n",
    "t_ex = np.arange(0, 60, 1/250)\n",
    "x_ex = np.sin(2*np.pi*f_ex*t_ex)\n",
    "peak_idxs_ex, peaks_ex = get_signal_extrema(x_ex, maxima=True)  \n",
    "assert np.all(np.isclose(ins_freqs(None, peak_idxs_ex)[1], f_ex, atol=0.01))\n",
    "plt.plot(np.arange(0, x_ex.shape[0]/250, 1/250), x_ex)\n",
    "plt.plot(peak_idxs_ex/250, x_ex[peak_idxs_ex], 'rx')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make Figure 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from plotly.subplots import make_subplots\n",
    "\n",
    "for idx in tqdm(range(len(marsh_dataset))):\n",
    "    if idx != 6:\n",
    "        continue\n",
    "        \n",
    "    for minute in range(3):\n",
    "        fp = marsh_dataset[idx].ECG().filepath.split('/')[-2]\n",
    "\n",
    "        # if remove when done\n",
    "        \n",
    "        print(f\"========================================================= Idx: {idx}, Subject: {fp}, minute: {minute+1} =========================================================\")\n",
    "        if (fp, minute) in exclusion_list:\n",
    "            print(\"SKIPPING\")\n",
    "            continue\n",
    "\n",
    "        # ------------------------------------------ METHOD A ------------------------------------------\n",
    "        # COMPARE WITH PEARSON\n",
    "        x_A = marsh_dataset[idx].ECG_ENV().transform(transforms=method_A_transforms[minute])\n",
    "        y = marsh_dataset[idx].IP().transform(transforms=gt_transforms[minute])\n",
    "        # convert PAC signal to usable format, then run CWT for WPC\n",
    "        try:\n",
    "            match = np.where(\n",
    "                (PAC_ids == np.array([np.where(PAC_subjects == fp)[0].item(), minute])).sum(axis=1) == 2\n",
    "            )[0].item()\n",
    "            \n",
    "            candB = PAC_signals[match]\n",
    "            x_B = Signal(_type='IP', data=candB, format='mat', filepath=None)\n",
    "            x_B = x_B.transform(transforms=method_A_transforms[minute] + [Transforms.MinMaxScale()]).transformed_data[200:]\n",
    "            x_A = x_A.transform(transforms=method_A_transforms[minute] + [Transforms.MinMaxScale()]).transformed_data[:x_B.shape[0]]\n",
    "            y = y.transform(transforms=method_A_transforms[minute] + [Transforms.MinMaxScale()]).transformed_data[:x_B.shape[0]]\n",
    "        \n",
    "            fig = make_subplots(rows=3, cols=1)\n",
    "            fig.add_trace(\n",
    "                go.Scatter(\n",
    "                    x = 60+np.arange(0, x_A.shape[0], 1/250),\n",
    "                    y = x_A,\n",
    "                    marker=dict(color='rgb(195,131,74)'),\n",
    "                    line=dict(width=2),\n",
    "                    name='ECG Envelope Signal',\n",
    "                ), row=1, col=1\n",
    "            )\n",
    "            fig.add_trace(\n",
    "                go.Scatter(\n",
    "                    x = 60+np.arange(0, x_B.shape[0], 1/250),\n",
    "                    y = x_B,\n",
    "                    marker=dict(color='rgb(73,131,73)'),\n",
    "                    line=dict(width=2),\n",
    "                    name='PAC Signal'\n",
    "                ), row=2, col=1\n",
    "            )\n",
    "            \n",
    "            fig.add_trace(\n",
    "                go.Scatter(\n",
    "                    x = 60+np.arange(0, y.shape[0], 1/250),\n",
    "                    y = y,\n",
    "                    # marker=dict(color='rgb(68,112,143)'), #original\n",
    "                    marker=dict(color='rgb(124, 218, 252)'),\n",
    "                    line=dict(width=2),\n",
    "                    name='IP Signal'\n",
    "                ), row=3, col=1\n",
    "            )\n",
    "            fig.update_layout(showlegend=False,plot_bgcolor='white', width=600, height=600, margin=dict(l=50, r=50, t=50, b=50), font_family='Arial',)\n",
    "            fig.update_xaxes(showline=True, linewidth=1, linecolor='black')\n",
    "            fig.update_yaxes(showline=True, linewidth=1, linecolor='black')\n",
    "            fig.update_xaxes(title_text='Time (s)', row=3, col=1)\n",
    "            fig.update_yaxes(title_text=\"Normalized Signal Amplitude\", row=2, col=1)\n",
    "            \n",
    "            fig.show()\n",
    "\n",
    "        except ValueError:\n",
    "            pass\n",
    "\n",
    "fig.write_image(\"paper_res/fig2.svg\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig.write_image(\"paper_res/fig2.svg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pipeline COmparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "most_prominent_freqs_A = []\n",
    "most_prominent_freqs_B = []\n",
    "# normalized phase coherences\n",
    "WPC_scores_A = []\n",
    "WPC_scores_B = []\n",
    "# non normalized\n",
    "WPCs_A = []\n",
    "WPCs_B = []\n",
    "\n",
    "# mean instantaneous frequencies, ground truth\n",
    "mean_ifs_gt_A = []\n",
    "mean_ifs_gt_B = []\n",
    "# mean IF's of estimated signal\n",
    "mean_ifs_est_A = []\n",
    "mean_ifs_est_B = []\n",
    "\n",
    "for idx in tqdm(range(len(marsh_dataset))):\n",
    "    for minute in range(3):\n",
    "        fp = marsh_dataset[idx].ECG().filepath.split('/')[-2]\n",
    "\n",
    "        # if remove when done\n",
    "        \n",
    "        print(f\"========================================================= Subject: {fp}, minute: {minute+1} =========================================================\")\n",
    "        if (fp, minute) in exclusion_list:\n",
    "            print(\"SKIPPING\")\n",
    "            continue\n",
    "\n",
    "        # ------------------------------------------ METHOD A ------------------------------------------\n",
    "        # COMPARE WITH WAVELET PHASE COHERENCE\n",
    "        candA = marsh_dataset[idx].ECG_ENV().transform(transforms=method_A_transforms[minute] + [CWT]).transformed_data\n",
    "        gtA = marsh_dataset[idx].IP().transform(transforms=gt_transforms[minute] + [CWT]).transformed_data\n",
    "        \n",
    "        t_A, PC_norms_A, PC_A = WPC(candA, gtA, freq=freq_band)\n",
    "        \n",
    "        avg_wpc_per_freq_A = PC_norms_A.mean(axis=1)\n",
    "        most_prominent_freqs_A.append(\n",
    "            freq_band[avg_wpc_per_freq_A.argmax()]\n",
    "        )\n",
    "        WPC_scores_A.append(\n",
    "            PC_norms_A.mean() # is a single value, averaging over time and frequency\n",
    "        )\n",
    "        WPCs_A.append(\n",
    "            PC_A.mean()\n",
    "        )\n",
    "        print(\"Method A freq, WPC norm, WPC not-norm\", most_prominent_freqs_A[-1], WPC_scores_A[-1], WPCs_A[-1])\n",
    "\n",
    "        # COMPARE WITH PEARSON\n",
    "        x_A = marsh_dataset[idx].ECG_ENV().transform(transforms=method_A_transforms[minute]).transformed_data\n",
    "        peak_idxs_A, peaks_A = get_signal_extrema(x_A, maxima=False)  \n",
    "        \n",
    "        y = marsh_dataset[idx].IP().transform(transforms=gt_transforms[minute]).transformed_data\n",
    "        annotated_peaks = marsh_dataset[idx].IP_annot().data \n",
    "        cropped_peaks = annotated_peaks[\n",
    "            (\n",
    "                annotated_peaks/250 > gt_transforms[minute][0].start\n",
    "            ) & (\n",
    "                annotated_peaks/250 < gt_transforms[minute][0].end\n",
    "            )\n",
    "        ] - gt_transforms[minute][0].start*250 # adjust indices to reference the start of the transformed data, not the raw data\n",
    "\n",
    "        mean_if_gt_A, mean_if_est_A = eval(*ins_freqs(y, cropped_peaks), *ins_freqs(x_A, peak_idxs_A))\n",
    "        mean_ifs_gt_A.append(mean_if_gt_A)\n",
    "        mean_ifs_est_A.append(mean_if_est_A)\n",
    "        print(f\"Mean IF of IP: {mean_ifs_gt_A[-1]}, Mean IF of estimate: {mean_ifs_est_A[-1]}\")\n",
    "\n",
    "        # ------------------------------------------ METHOD B ------------------------------------------\n",
    "        try:\n",
    "            match = np.where(\n",
    "                (PAC_ids == np.array([np.where(PAC_subjects == fp)[0].item(), minute])).sum(axis=1) == 2\n",
    "            )[0].item()\n",
    "            \n",
    "            candB = PAC_signals[match]\n",
    "            # convert PAC signal to usable format, then run CWT for WPC\n",
    "            x_B = Signal(_type='IP', data=candB, format='mat', filepath=None)\n",
    "            t_B, PC_norms_B, PC_B = WPC(CWT(x_B), gtA, freq=freq_band)\n",
    "\n",
    "            avg_wpc_per_freq_B = PC_norms_B.mean(axis=1)\n",
    "            \n",
    "            most_prominent_freqs_B.append(\n",
    "                freq_band[avg_wpc_per_freq_B.argmax()]\n",
    "            )\n",
    "            WPC_scores_B.append(\n",
    "                PC_norms_B.mean() # is a single value, averaging over time and frequency\n",
    "            )\n",
    "            WPCs_B.append(\n",
    "                PC_B.mean()\n",
    "            )\n",
    "            print(\"Method B freq, WPC norm, WPC not-norm\", most_prominent_freqs_B[-1], WPC_scores_B[-1], WPCs_B[-1])\n",
    "    \n",
    "            # COMPARE WITH PEARSON\n",
    "            peak_idxs_B, peaks_B = get_signal_extrema(candB, maxima=True)  \n",
    "            \n",
    "            mean_if_gt_B, mean_if_est_B = eval(*ins_freqs(y, cropped_peaks), *ins_freqs(candB, peak_idxs_B))\n",
    "            mean_ifs_gt_B.append(mean_if_gt_B)\n",
    "            mean_ifs_est_B.append(mean_if_est_B)\n",
    "            print(f\"Mean IF of IP: {mean_ifs_gt_B[-1]}, Mean IF of estimate: {mean_ifs_est_B[-1]}\")\n",
    "        \n",
    "        except ValueError:\n",
    "            print(\"PAC Match not found\")\n",
    "                \n",
    "            most_prominent_freqs_B.append(None)\n",
    "            WPC_scores_B.append(None)\n",
    "            WPCs_B.append(None)\n",
    "            mean_ifs_gt_B.append(None)\n",
    "            mean_ifs_est_B.append(None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot minute 7, guided breathing at 0.133hz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metronome_7_min_transforms = [Transforms.Crop(start=60*7, end=60*(7+1))] + pipeline\n",
    "marsh_dataset[idx].IP().transform(transforms=metronome_7_min_transforms).plot(transformed=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "marsh_dataset[idx].ECG_ENV().transform(transforms=metronome_7_min_transforms).plot(transformed=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "idx=4\n",
    "candA = marsh_dataset[idx].ECG_ENV().transform(transforms=metronome_7_min_transforms + [CWT]).transformed_data\n",
    "gtA = marsh_dataset[idx].IP().transform(transforms=metronome_7_min_transforms + [CWT]).transformed_data\n",
    "t_A, PC_norms_A, PC_A = WPC(candA, gtA, freq=freq_band)\n",
    "plt.figure(figsize=(10, 6))\n",
    "sns.heatmap(PC_norms_A, xticklabels=t_A.round(2), yticklabels=freq_band.round(2))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_fig = marsh_dataset[idx].ECG_ENV().transform(transforms=method_A_transforms[1] + [lambda x: -1*x]).plot(transformed=True)\n",
    "base_fig = marsh_dataset[idx].IP().transform(transforms=method_A_transforms[1]).plot(transformed=True, base_fig=base_fig)\n",
    "base_fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "subject_ids = []\n",
    "minutes = []\n",
    "for idx in tqdm(range(len(marsh_dataset))):\n",
    "    for minute in range(3):\n",
    "        fp = marsh_dataset[idx].ECG().filepath.split('/')[-2]\n",
    "        if (fp, minute) in exclusion_list:\n",
    "            continue\n",
    "        subject_ids.append(fp)\n",
    "        minutes.append(minute)\n",
    "        \n",
    "results_df = pd.DataFrame({\n",
    "    \"subject_id\": subject_ids,\n",
    "    \"minute\": minutes,\n",
    "    \"WPC_scores_A\": WPC_scores_A,\n",
    "    \"WPC_scores_B\": WPC_scores_B,\n",
    "    \"WPCs_A_non_normalized\": WPCs_A,\n",
    "    \"WPCs_B_non_normalized\": WPCs_B,\n",
    "    \"mean_ifs_gt_A\": mean_ifs_gt_A,\n",
    "    \"mean_ifs_est_A\": mean_ifs_est_A,\n",
    "    \"mean_ifs_gt_B\": mean_ifs_gt_B,\n",
    "    \"mean_ifs_est_B\": mean_ifs_est_B,\n",
    "})\n",
    "display(results_df)\n",
    "results_df.to_csv(f\"PAC_ECG_ENV_results_lo_{lo}hz_hi_{hi}hz.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_results_df = results_df[~results_df.WPC_scores_B.isna()]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparing Results By Minute"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(results_df.groupby(by='minute').WPCs_A_non_normalized.mean())\n",
    "display(results_df.groupby(by='minute').WPCs_B_non_normalized.mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame({'Ground Truth':mean_ifs_gt_A, 'ECG Envelope':mean_ifs_est_A, \"PAC\": mean_ifs_est_B})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"METHOD A RESULTS\")\n",
    "WPC_avg_A, WPC_std_A = filtered_results_df.WPCs_A_non_normalized.mean(), filtered_results_df.WPCs_A_non_normalized.std()\n",
    "print(\"WPC Score Avg and Std\", WPC_avg_A, WPC_std_A)\n",
    "mean_ifs_gt_A, mean_ifs_est_A = filtered_results_df.mean_ifs_gt_A.to_numpy(), filtered_results_df.mean_ifs_est_A.to_numpy()\n",
    "plt.plot([0, 1], [0, 1], ls=\"--\")\n",
    "plt.plot(mean_ifs_gt_A, mean_ifs_est_A, 'bo')\n",
    "plt.xlabel(\"Ground Truth Inst. Freqs\")\n",
    "plt.ylabel(\"Extracted Signal Inst. Freqs\")\n",
    "plt.xlim([0.05, 0.35])\n",
    "plt.ylim([0.05, 0.35])\n",
    "plt.show()\n",
    "PCC_A = scipy.stats.pearsonr(mean_ifs_gt_A, mean_ifs_est_A)[0]\n",
    "print(\"PCC A:\", PCC_A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"METHOD B RESULTS\")\n",
    "most_prominent_freqs_B, WPC_scores_B = np.array(most_prominent_freqs_B), np.array(WPC_scores_B)\n",
    "WPC_avg_B, WPC_std_B = filtered_results_df.WPCs_B_non_normalized.mean(), filtered_results_df.WPCs_B_non_normalized.std()\n",
    "print(\"WPC Score Avg and Std:\", WPC_avg_B, WPC_std_B)\n",
    "mean_ifs_gt_B, mean_ifs_est_B = filtered_results_df.mean_ifs_gt_B.to_numpy(), filtered_results_df.mean_ifs_est_B.to_numpy()\n",
    "plt.plot([0, 1], [0, 1], ls=\"--\")\n",
    "plt.plot(mean_ifs_gt_B[mean_ifs_gt_B.astype(bool)], mean_ifs_est_B[mean_ifs_est_B.astype(bool)], 'bo')\n",
    "plt.xlabel(\"Ground Truth Inst. Freqs\")\n",
    "plt.ylabel(\"Extracted Signal Inst. Freqs\")\n",
    "plt.xlim([0.05, 0.35])\n",
    "plt.ylim([0.05, 0.35])\n",
    "plt.show()\n",
    "PCC_B = scipy.stats.pearsonr(mean_ifs_gt_B[mean_ifs_gt_B.astype(bool)], mean_ifs_est_B[mean_ifs_est_B.astype(bool)])[0]\n",
    "print(\"PCC B:\", PCC_B)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_ifs_gt_B[mean_ifs_gt_B.astype(bool)].mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"{PCC_A:.3f} & {PCC_B:.3f} & {WPC_avg_A:.3f} & {WPC_avg_B:.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from plotly.subplots import make_subplots\n",
    "import plotly.graph_objects as go\n",
    "\n",
    "fig = make_subplots(rows=1, cols=1, subplot_titles=[''])\n",
    "\n",
    "fig.add_trace(go.Bar(\n",
    "    x=['ECG Envelope'],\n",
    "    y=[WPC_avg_A],\n",
    "    marker=dict(color='rgb(195,131,74)'),\n",
    "    error_y=dict(\n",
    "            type='data', # value of error bar given in data coordinates\n",
    "            array=[WPC_std_A],\n",
    "            visible=True\n",
    "    )\n",
    "), row=1, col=1)\n",
    "fig.add_trace(go.Bar(\n",
    "    x=['PAC'],\n",
    "    y=[WPC_avg_B],\n",
    "    marker=dict(color='rgb(73,131,73)'),\n",
    "    error_y=dict(\n",
    "            type='data', # value of error bar given in data coordinates\n",
    "            array=[WPC_std_B],\n",
    "            visible=True\n",
    "    )\n",
    "), row=1, col=1)\n",
    "\n",
    "fig.update_layout(plot_bgcolor='white', height=450, width=500, showlegend=False, font_family='Arial', yaxis_title='Average WPC')\n",
    "fig.update_xaxes(showline=True, linewidth=1, linecolor='black')\n",
    "fig.update_yaxes(showline=True, linewidth=1, linecolor='black', range=[0.75, 0.975])\n",
    "fig.show()\n",
    "\n",
    "# plt.bar(\n",
    "#     ['ECG Envelope', 'PAC'],\n",
    "#     [WPC_avg_A, WPC_avg_B],\n",
    "# )\n",
    "# plt.errorbar(\n",
    "#     ['ECG Envelope', 'PAC'],\n",
    "#     [WPC_avg_A, WPC_avg_B],\n",
    "#     [WPC_std_A, WPC_std_B],\n",
    "#     fmt='.', color='Black', elinewidth=2,capthick=10,errorevery=1, alpha=0.5, ms=4, capsize = 2\n",
    "# )\n",
    "# plt.ylabel(\"Average Wavelet Phase Coherence\")\n",
    "# plt.ylim([0.7, 1])\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.bar(\n",
    "    ['ECG Envelope', 'PAC'],\n",
    "    [PCC_A, PCC_B],\n",
    ")\n",
    "plt.ylabel(\"Pearson Correlation Coefficient\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Paired T-test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta_df = pd.DataFrame({\n",
    "    \"WPC ECG ENV\": filtered_results_df.WPCs_A_non_normalized, \n",
    "    \"WPC PAC\" : filtered_results_df.WPCs_B_non_normalized,\n",
    "    \"Delta\": filtered_results_df.WPCs_A_non_normalized-filtered_results_df.WPCs_B_non_normalized\n",
    "})\n",
    "\n",
    "delta_df.sample(20).sort_values(by='Delta', ascending=False).style.background_gradient(cmap=\"RdYlGn\", subset=['Delta'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Manual Result\")\n",
    "print(delta_df.Delta.mean()/(\n",
    "    delta_df.Delta.std()/np.sqrt(len(delta_df))\n",
    "))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import ttest_rel\n",
    "# print(\n",
    "#     ttest_rel(filtered_results_df.WPCs_A_non_normalized, filtered_results_df.WPCs_B_non_normalized, alternative='two-sided')\n",
    "# )\n",
    "print(\"Scipy Result\")\n",
    "print(\n",
    "    ttest_rel(filtered_results_df.WPCs_A_non_normalized, filtered_results_df.WPCs_B_non_normalized, alternative='greater').statistic\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Figure Comparing signals produced by both methods vs Ground Truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for idx in tqdm(range(len(marsh_dataset))):\n",
    "    for minute in range(3):\n",
    "        fp = marsh_dataset[idx].ECG().filepath.split('/')[-2]\n",
    "\n",
    "        # if remove when done\n",
    "        \n",
    "        print(f\"========================================================= Subject: {fp}, minute: {minute+1} dataset idx: {idx} =========================================================\")\n",
    "        if (fp, minute) in exclusion_list:\n",
    "            print(\"SKIPPING\")\n",
    "            continue\n",
    "\n",
    "        # ------------------------------------------ METHOD A ------------------------------------------\n",
    "\n",
    "        x_A = marsh_dataset[idx].ECG_ENV().transform(transforms=method_A_transforms[minute]).transformed_data\n",
    "        \n",
    "        y = marsh_dataset[idx].IP().transform(transforms=gt_transforms[minute]).transformed_data\n",
    "\n",
    "        # ------------------------------------------ METHOD B ------------------------------------------\n",
    "        try:\n",
    "            match = np.where(\n",
    "                (PAC_ids == np.array([np.where(PAC_subjects == fp)[0].item(), minute])).sum(axis=1) == 2\n",
    "            )[0].item()\n",
    "            \n",
    "            candB = PAC_signals[match]\n",
    "            # convert PAC signal to usable format, then run CWT for WPC\n",
    "            x_B = Signal(_type='IP', data=candB, format='mat', filepath=None)\n",
    "                        \n",
    "            # Create a figure and a set of subplots\n",
    "            fig, axs = plt.subplots(3, 1, figsize=(8, 5))  # width = 8 inches, height = 5 inches\n",
    "            \n",
    "            # Plot for Method A\n",
    "            axs[0].plot(\n",
    "                np.arange(0, x_A.shape[0]/250, 1/250),\n",
    "                Transforms.MinMaxScale()(x_A),\n",
    "                label='ECG Envelope'\n",
    "            )\n",
    "            axs[0].set_title('ECG Envelope')\n",
    "            axs[0].grid(True)\n",
    "            \n",
    "            # Plot for Method B\n",
    "            axs[1].plot(\n",
    "                np.arange(0, candB.shape[0]/250, 1/250),\n",
    "                Transforms.MinMaxScale()(candB),\n",
    "                label='PAC'\n",
    "            )\n",
    "            axs[1].set_title('PAC')\n",
    "            axs[1].grid(True)\n",
    "            \n",
    "            # Plot for True Signal\n",
    "            axs[2].plot(\n",
    "                np.arange(0, y.shape[0]/250, 1/250), \n",
    "                Transforms.MinMaxScale()(y), \n",
    "                label='True Signal'\n",
    "            )\n",
    "            axs[2].set_title('True Signal')\n",
    "            axs[2].grid(True)\n",
    "            \n",
    "            # Setting the background color of the figure\n",
    "            fig.patch.set_facecolor('white')\n",
    "            \n",
    "            # Setting the spines for each axis\n",
    "            for ax in axs:\n",
    "                for spine in ax.spines.values():\n",
    "                    spine.set_edgecolor('black')\n",
    "                    spine.set_linewidth(1)\n",
    "            \n",
    "            # Adjusting layout\n",
    "            plt.tight_layout()\n",
    "            \n",
    "            # Show the plot\n",
    "            plt.show()\n",
    "            \n",
    "        except ValueError as e:\n",
    "            print(\"PAC Match not found\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from plotly.subplots import make_subplots\n",
    "# idx = 3\n",
    "# minute = 1\n",
    "# fp = marsh_dataset[idx].ECG().filepath.split('/')[-2]\n",
    "\n",
    "# print(f\"======================================================= Subject: {fp}, minute: {minute+1} dataset idx: {idx} =======================================================\")\n",
    "\n",
    "# # ------------------------------------------ METHOD A ------------------------------------------\n",
    "\n",
    "# x_A = marsh_dataset[idx].ECG_ENV().transform(transforms=method_A_transforms[minute]).transformed_data\n",
    "\n",
    "# y = marsh_dataset[idx].IP().transform(transforms=gt_transforms[minute]).transformed_data\n",
    "\n",
    "# # ------------------------------------------ METHOD B ------------------------------------------\n",
    "# match = np.where(\n",
    "#     (PAC_ids == np.array([np.where(PAC_subjects == fp)[0].item(), minute])).sum(axis=1) == 2\n",
    "# )[0].item()\n",
    "\n",
    "# candB = PAC_signals[match]\n",
    "# # convert PAC signal to usable format, then run CWT for WPC\n",
    "# x_B = Signal(_type='IP', data=candB, format='mat', filepath=None)\n",
    "\n",
    "fig = make_subplots(rows=3, cols=1)\n",
    "fig.add_trace(\n",
    "    go.Scatter(\n",
    "        x = 60+np.arange(0, x_A.shape[0], 1/250),\n",
    "        y = Transforms.MinMaxScale()(x_A),\n",
    "        marker=dict(color='rgb(195,131,74)'),\n",
    "        line=dict(width=5),\n",
    "        name='ECG Envelope Signal'\n",
    "    ), row=1, col=1\n",
    ")\n",
    "fig.add_trace(\n",
    "    go.Scatter(\n",
    "        x = 60+np.arange(0, x_B.data.shape[0], 1/250),\n",
    "        y = Transforms.MinMaxScale()(x_B.data),\n",
    "        marker=dict(color='rgb(73,131,73)'),\n",
    "        line=dict(width=5),\n",
    "        name='PAC Signal'\n",
    "    ), row=2, col=1\n",
    ")\n",
    "\n",
    "fig.add_trace(\n",
    "    go.Scatter(\n",
    "        x = 60+np.arange(0, y.shape[0], 1/250),\n",
    "        y = Transforms.MinMaxScale()(y),\n",
    "        # marker=dict(color='rgb(68,112,143)'), #original\n",
    "        marker=dict(color='rgb(124, 218, 252)'),\n",
    "        line=dict(width=5),\n",
    "        name='IP Signal'\n",
    "    ), row=3, col=1\n",
    ")\n",
    "fig.update_layout(plot_bgcolor='white', width=800, height=500)\n",
    "fig.update_xaxes(showline=True, linewidth=1, linecolor='black')\n",
    "fig.update_yaxes(showline=True, linewidth=1, linecolor='black')\n",
    "fig.update_xaxes(title_text='Time (s)', row=3, col=1)\n",
    "fig.update_yaxes(title_text=\"Normalized Signal Amplitude\", row=2, col=1)\n",
    "\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exporting ECG ENV data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# exported_ecg_envs = []\n",
    "# exported_fps = []\n",
    "# exported_minutes = []\n",
    "\n",
    "\n",
    "# for idx in tqdm(range(len(marsh_dataset))):\n",
    "#     for minute in range(3):\n",
    "\n",
    "#         fp = marsh_dataset[idx].ECG().filepath.split('/')[-2]\n",
    "\n",
    "#         # if remove when done\n",
    "        \n",
    "#         print(f\"========================================================= Subject: {fp}, minute: {minute+1} =========================================================\")\n",
    "#         if (fp, minute) in exclusion_list:\n",
    "#             print(\"SKIPPING\")\n",
    "#             continue\n",
    "\n",
    "#         exported_ecg_envs.append(marsh_dataset[idx].ECG_ENV().transform(transforms=method_A_transforms[minute]).transformed_data)\n",
    "#         exported_fps.append(fp)\n",
    "#         exported_minutes.append(minute)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from scipy.io import savemat\n",
    "# export_data = {\n",
    "#     \"ECG_ENV\": exported_ecg_envs,\n",
    "#     \"Subject_ID\": exported_fps,\n",
    "#     \"Minute\": exported_minutes\n",
    "# }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# savemat(\"marsh_ECG_envelopes_jan_2024.mat\", export_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Eval Functions (not for this paper)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = pd.DataFrame({\n",
    "    'Instantaneous_Frequency_CorrA': WPC_scores_A,\n",
    "    'Instantaneous_Frequency_CorrB': WPC_scores_B,\n",
    "    'Coherence_A': cohsA,\n",
    "    'Coherence_B': cohsB,\n",
    "})\n",
    "df['dataset'] = ['MARSH']*29 + ['Fantasia']*28"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from plotly.subplots import make_subplots\n",
    "import plotly.graph_objects as go\n",
    "\n",
    "fig = make_subplots(rows=1, cols=2, subplot_titles=['Ins. Freq Correlation', 'Avg Band Coherence'])\n",
    "\n",
    "mean_df = df.groupby(by='dataset').mean()\n",
    "fig.add_trace(go.Bar(\n",
    "    x = [\"Fantasia\", \"Marsh\"],# \"COHa\", \"COHb\"],\n",
    "    y = mean_df.iloc[:, 0].T, \n",
    "    name='Method A (ECG Env)',\n",
    "    marker=dict(color='blue')\n",
    "), row=1, col=1)\n",
    "fig.add_trace(go.Bar(\n",
    "    x = [\"Fantasia\", \"Marsh\"],# \"COHa\", \"COHb\"],\n",
    "    y = mean_df.iloc[:, 1], \n",
    "    name='Method B (IMFs)',\n",
    "    marker=dict(color='red')\n",
    "), row=1, col=1)\n",
    "fig.add_trace(go.Bar(\n",
    "    x = [\"Fantasia\", \"Marsh\"],# \"COHa\", \"COHb\"],\n",
    "    y = mean_df.iloc[:, 2].T, \n",
    "    name='Method A (ECG Env)',\n",
    "    marker=dict(color='blue')\n",
    "), row=1, col=2)\n",
    "fig.add_trace(go.Bar(\n",
    "    x = [\"Fantasia\", \"Marsh\"],# \"COHa\", \"COHb\"],\n",
    "    y = mean_df.iloc[:, 3], \n",
    "    name='Method B (IMFs)',\n",
    "    marker=dict(color='red')\n",
    "), row=1, col=2)\n",
    "fig.update_layout(plot_bgcolor='white', width=800)\n",
    "fig.update_xaxes(showline=True, linewidth=1, linecolor='black')\n",
    "fig.update_yaxes(showline=True, linewidth=1, linecolor='black')\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def choose_best_imf(imfs, L_br=1/16, U_br=1/8, imf_bar_plot=False, **kwargs):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "    - imfs:\n",
    "    - n_candidates: \n",
    "    - L_br, U_br: lower and upper bound of respiratory rate, will be used to extract the IMFs with the most \n",
    "    power in this band. Extracted from bounds established in\n",
    "    https://www.ncbi.nlm.nih.gov/pmc/articles/PMC7796076/#:~:text=Due%20to%20the%20displacement%20of,by%20body%20movement%2C%20including%20breathing.\n",
    "    \"\"\"\n",
    "    w = scipy.signal.windows.tukey(imfs.shape[1], alpha=0.1)\n",
    "    # proportion of power that is in the range of breathing\n",
    "    band_powers = []\n",
    "    for i, imf in enumerate(imfs):\n",
    "        \n",
    "        imf_lp = butter_lowpass_filter(w*imf, cutoff=40, fs=250, order=4)\n",
    "        freqs_x, fft_x = FFT(imf_lp, cutoff=5)\n",
    "        band_powers.append(\n",
    "            fft_x[(freqs_x > L_br) & (freqs_x < U_br)].sum()/fft_x.sum()\n",
    "        )\n",
    "    if imf_bar_plot:\n",
    "        plt.title(\"IMF Band Powers\")\n",
    "        plt.bar(np.arange(imfs.shape[0]), band_powers)\n",
    "        plt.show()\n",
    "        \n",
    "    return imfs[np.argmax(band_powers)]\n",
    "\n",
    "def plot_imfs(imfs, ip=None):\n",
    "    fig, ax = plt.subplots(imfs.shape[0], 2, figsize=(10, imfs.shape[0]*2.5))\n",
    "    for i, imf in enumerate(imfs):\n",
    "        freqs_x, fft_x = FFT(imf, cutoff=5)\n",
    "        \n",
    "        ax[i, 0].set_title(f\"IMF {i}\")\n",
    "        ax[i, 0].plot(np.arange(0, imf.shape[0]/250, 1/250), imf)\n",
    "        ax[i, 1].set_title(f\"Energy of Br Band: {fft_x[(freqs_x > L_br) & (freqs_x < U_br)].sum()/fft_x.sum()}\")\n",
    "        ax[i, 1].plot(freqs_x, fft_x)\n",
    "        \n",
    "        if ip is not None:\n",
    "            freqs_ip, fft_ip = FFT(ip, cutoff=5)\n",
    "            ax[i, 0].plot(np.arange(0, ip.shape[0]/250, 1/250), ip)\n",
    "            ax[i, 1].plot(freqs_ip, fft_ip)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "def MethodB(signal, ecg_transforms, n_candidates=2, trials=100, noise_width=0.8, **kwargs):\n",
    "    x = signal.ECG().transform(transforms=ecg_transforms).transformed_data\n",
    "    L_br, U_br = 1/16, 0.3\n",
    "    # remove linear trend and DC\n",
    "    x -= x.mean()\n",
    "    x = detrend(x)\n",
    "    # calculate IMFs ()\n",
    "    imfs =  EEMD(trials=trials, noise_width=noise_width)(x)\n",
    "    # choose best IMF()\n",
    "    candidate = choose_best_imf(imfs, **kwargs)\n",
    "    # return time, and best predictions\n",
    "    return candidate, imfs\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Run On Both Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "crop = Transforms.Crop(start=5, end=180)\n",
    "resultsA = []\n",
    "cohsA = []\n",
    "resultsB = []\n",
    "cohsB = []\n",
    "for i in tqdm(range(len(marsh_dataset))):\n",
    "    print(f\"---------------------------- 'MARSH: SAMPLE {i}--------------------------------------------------------\")\n",
    "    ground_truth = marsh_dataset[i].IP().transform(transforms=[\n",
    "        crop,  \n",
    "        lambda x: x-x.mean(),\n",
    "        lambda x: detrend(x),\n",
    "        lambda x: np.convolve(np.ones(250)/250, x, mode='same') # smooth\n",
    "    ]).transformed_data\n",
    "    \n",
    "    candA = marsh_dataset[i].ECG_ENV().transform(transforms=[\n",
    "        crop,  \n",
    "        lambda x: x-x.mean(),\n",
    "        lambda x: detrend(x),\n",
    "    ]).transformed_data\n",
    "    f, CxyA = scipy.signal.coherence(candA, ground_truth)\n",
    "    \n",
    "    resultsA.append(eval(ground_truth=ground_truth, candidate=candA))\n",
    "    cohsA.append(CxyA[(f > 1/16) & (f < 0.3)].mean())\n",
    "    print(\n",
    "        f\"Method A: {resultsA[-1]} coherence {cohsA[-1]}\"\n",
    "    )\n",
    "    \n",
    "    # COMMENTED OUT IMF STUFF FOR LATER\n",
    "    # candB, imfs = MethodB(\n",
    "    #     marsh_dataset[i],\n",
    "    #     ecg_transforms=[crop],\n",
    "    #     trials=30,\n",
    "    #     # imf_bar_plot=True\n",
    "    # )\n",
    "    # f, CxyB = scipy.signal.coherence(candB, ground_truth)\n",
    "    \n",
    "    # resultsB.append(eval(ground_truth=ground_truth, candidate=candB))\n",
    "    # cohsB.append(CxyB[(f > 1/16) & (f < 0.3)].mean())\n",
    "    # print(\n",
    "    #     f\"Method B: {resultsB[-1]} coherence {cohsB[-1]}\"\n",
    "    # )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for i in tqdm(range(len(fantasia_dataset))):\n",
    "    for j in range(7):\n",
    "        crop = Transforms.Crop(start=(180*j)+5, end=180*(j+1))\n",
    "        \n",
    "        print(f\"---------------------------- Fantasia: SAMPLE {i}, SEGMENT {j}--------------------------------------------------------\")\n",
    "        ground_truth = fantasia_dataset[i].IP().transform(transforms=[\n",
    "            crop,  \n",
    "            lambda x: x-x.mean(),\n",
    "            lambda x: detrend(x),\n",
    "            lambda x: np.convolve(np.ones(250)/250, x, mode='same') # smooth\n",
    "        ]).transformed_data\n",
    "        \n",
    "        candA = fantasia_dataset[i].ECG_ENV().transform(transforms=[\n",
    "            crop,\n",
    "            lambda x: x-x.mean(),\n",
    "            lambda x: detrend(x),\n",
    "        ]).transformed_data\n",
    "        f, CxyA = scipy.signal.coherence(candA, ground_truth)\n",
    "        \n",
    "        resultsA.append(eval(ground_truth=ground_truth, candidate=candA))\n",
    "        cohsA.append(CxyA[(f > 1/16) & (f < 0.3)].mean())\n",
    "        print(\n",
    "            f\"Method A: {resultsA[-1]} coherence {cohsA[-1]}\"\n",
    "        )\n",
    "\n",
    "        # IMF STUFF COMMENTED OUT\n",
    "        # candB, imfs = MethodB(\n",
    "        #     fantasia_dataset[i],\n",
    "        #     ecg_transforms=[crop],\n",
    "        #     trials=30,\n",
    "        #     # imf_bar_plot=True\n",
    "        # )\n",
    "        # f, CxyB = scipy.signal.coherence(candB, ground_truth)\n",
    "        \n",
    "        # resultsB.append(eval(ground_truth=ground_truth, candidate=candB))\n",
    "        # cohsB.append(CxyB[(f > 1/16) & (f < 0.3)].mean())\n",
    "        # print(\n",
    "        #     f\"Method B: {resultsB[-1]} coherence {cohsB[-1]}\"\n",
    "        # )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# XGB Regressor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xgboost as xgb\n",
    "import scipy\n",
    "import matplotlib.pyplot as plt\n",
    "from numpy.lib.stride_tricks import as_strided"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_rolling_windows(signal, window_size=1000):\n",
    "        \n",
    "    # Calculate new shape and strides\n",
    "    new_shape = (signal.size - window_size + 1, window_size)\n",
    "    new_strides = (signal.strides[0], signal.strides[0])\n",
    "    \n",
    "    # Create the rolling window view\n",
    "    return as_strided(signal, shape=new_shape, strides=new_strides)\n",
    "\n",
    "def get_masks(signal, p=0.08, mask_length=125):\n",
    "    return np.random.choice([0, 1], size=(signal.shape[0], mask_length), p=[1,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import roc_auc_score, confusion_matrix\n",
    "from tqdm.auto import tqdm\n",
    "    \n",
    "train_idx = 0\n",
    "start_time_XGB = 0 # in seconds\n",
    "end_time_XGB = 180 # in seconds\n",
    "\n",
    "N_test = 60 # one minute of data\n",
    "start_time_XGB_test = 420"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_fig = marsh_dataset[train_idx].ECG().transform(transforms=[\n",
    "        Transforms.Crop(start_time_XGB, end_time_XGB),\n",
    "        Transforms.MinMaxScale(),\n",
    "        lambda x: x-x.mean(),\n",
    "        lambda x: detrend(x),\n",
    "    ]).plot(transformed=True)\n",
    "\n",
    "\n",
    "marsh_dataset[train_idx].IP().transform(transforms=[\n",
    "    Transforms.Crop(start_time_XGB, end_time_XGB),\n",
    "    Transforms.MinMaxScale(),\n",
    "    lambda x: x-x.mean(),\n",
    "    lambda x: detrend(x),\n",
    "]).plot(transformed=True, base_fig=train_fig)\n",
    "\n",
    "train_fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test_idx = 8\n",
    "# # start_time_XGB = 240 # in seconds\n",
    "# # end_time_XGB = 700 # in seconds\n",
    "    \n",
    "# WPC_per_index = []\n",
    "# WPC_smoothed_per_index = []\n",
    "# for train_idx in tqdm(range(len(marsh_dataset))):#NUM_MARSH_TRAIN_SAMPLES_XGB)):\n",
    "\n",
    "#     print(\"================================================================================\")\n",
    "#     print(\"INDEX:\", train_idx)\n",
    "#     # ECG input data\n",
    "#     X_ecg_rolling_train_stack = []\n",
    "#     y_ip_train_stack = []\n",
    "    \n",
    "#     input_ecg_raw = marsh_dataset[train_idx].ECG().transform(transforms=[\n",
    "#         Transforms.Crop(start_time_XGB, end_time_XGB),\n",
    "#         Transforms.MinMaxScale(),\n",
    "#         lambda x: x-x.mean(),\n",
    "#         lambda x: detrend(x),\n",
    "#     ])\n",
    "    \n",
    "#     # Breathing Rate input and target data\n",
    "#     input_ip_raw = marsh_dataset[train_idx].IP().transform(transforms=[\n",
    "#         Transforms.Crop(start_time_XGB, end_time_XGB),\n",
    "#         Transforms.MinMaxScale(),\n",
    "#         lambda x: x-x.mean(),\n",
    "#         lambda x: detrend(x),\n",
    "#     ])\n",
    "    \n",
    "#     # rolling window parameters\n",
    "#     window_size = input_ecg_raw.sample_rate*4\n",
    "    \n",
    "#     input_ecg = input_ecg_raw.transformed_data\n",
    "#     input_ip = input_ip_raw.transformed_data\n",
    "    \n",
    "#     X_ecg_rolling = get_rolling_windows(input_ecg, window_size=window_size)\n",
    "#     y_ip = input_ip.data[window_size-1:]\n",
    "    \n",
    "#     X_ecg_rolling_train_stack.append(X_ecg_rolling)\n",
    "#     y_ip_train_stack.append(y_ip)\n",
    "\n",
    "#     X_ecg_rolling_train_stack_np = np.vstack(X_ecg_rolling_train_stack)\n",
    "#     y_ip_train_stack_np = np.stack(y_ip_train_stack).flatten()\n",
    "#     # print(X_ecg_rolling_train_stack_np.shape, y_ip_train_stack_np.shape)\n",
    "    \n",
    "#     reg = xgb.XGBRegressor(tree_method=\"hist\")\n",
    "#     # Fit the model using predictor X and response y.\n",
    "#     reg.fit(X_ecg_rolling_train_stack_np, y_ip_train_stack_np)\n",
    "\n",
    "\n",
    "#     ## Set up test data (from 420s to 480s, the controlled breathing at 0.133hz\n",
    "#     # ECG input data\n",
    "#     x_test_ecg_raw = marsh_dataset[train_idx].ECG().transform(transforms=[\n",
    "#         Transforms.Crop(start_time_XGB_test, start_time_XGB_test+N_test),\n",
    "#         Transforms.MinMaxScale(),\n",
    "#         lambda x: x-x.mean(),\n",
    "#         lambda x: detrend(x),\n",
    "#     ]).transformed_data\n",
    "    \n",
    "#     x_test_ecg_rolling = get_rolling_windows(x_test_ecg_raw, window_size=window_size)\n",
    "    \n",
    "#     y_test = marsh_dataset[train_idx].IP().transform(transforms=[\n",
    "#         Transforms.Crop(start_time_XGB_test, start_time_XGB_test+N_test),\n",
    "#         Transforms.MinMaxScale(),\n",
    "#         lambda x: x-x.mean(),\n",
    "#         lambda x: detrend(x),\n",
    "#     ]).transformed_data[window_size-1:]\n",
    "    \n",
    "#     preds = []\n",
    "#     smoothed_preds = []\n",
    "#     smoothing_length = 50\n",
    "#     mean_filter = np.ones(smoothing_length)\n",
    "    \n",
    "#     for xi in tqdm(x_test_ecg_rolling):\n",
    "#         pred = reg.predict(xi[np.newaxis, :])[0]\n",
    "#         preds.append(pred)\n",
    "#         filt_order = min(smoothing_length, len(preds))\n",
    "#         smoothed_preds.append(\n",
    "#             np.dot(preds[-filt_order:],mean_filter[-filt_order:])/filt_order # smoothed prediction\n",
    "#         )\n",
    "#     preds = np.array(preds)\n",
    "#     smoothed_preds = np.array(smoothed_preds)\n",
    "#     assert preds.shape[0] == y_test.shape[0]\n",
    "    \n",
    "#     fig = go.Figure()\n",
    "#     fig.add_trace(go.Scatter(y = preds, name='pred'))\n",
    "#     fig.add_trace(go.Scatter(y = smoothed_preds, name='smoothed pred'))\n",
    "#     fig.add_trace(go.Scatter(y = y_test, name='ground truth'))\n",
    "#     fig.show()\n",
    "    \n",
    "#     pred_signal = Signal(format='mat', filepath=None, _type='IP', data=np.array(preds))\n",
    "#     pred_smoothed_signal = Signal(format='mat', filepath=None, _type='IP', data=np.array(smoothed_preds))\n",
    "#     gt_signal = Signal(format='mat', filepath=None, _type='IP', data=np.array(y_test))\n",
    "\n",
    "#     pred_signal_CWT = pred_signal.transform(transforms=[CWT]).transformed_data\n",
    "#     pred_smoothed_signal_CWT = pred_smoothed_signal.transform(transforms=[CWT]).transformed_data\n",
    "#     gt_signal_CWT = gt_signal.transform(transforms=[CWT]).transformed_data\n",
    "#     _, _, PC_test = WPC(\n",
    "#         pred_signal_CWT,\n",
    "#         gt_signal_CWT,\n",
    "#     )\n",
    "#     _, _, PC_smoothed_test = WPC(\n",
    "#         pred_smoothed_signal_CWT,\n",
    "#         gt_signal_CWT,\n",
    "#     )\n",
    "    \n",
    "#     WPC_per_index.append(PC_test.mean())\n",
    "#     WPC_smoothed_per_index.append(PC_smoothed_test.mean())\n",
    "    \n",
    "#     print(f\"| WPC | Smoothed WPC |\")\n",
    "#     print(f\"| {WPC_per_index[-1]:.4f} | {WPC_smoothed_per_index[-1]:.4f} |\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fig = go.Figure()\n",
    "# fig.add_trace(\n",
    "#     go.Bar(\n",
    "#         x=np.array(range(len(WPC_per_index))), \n",
    "#         y=np.array(WPC_per_index),\n",
    "#         name=\"WPC: Raw Prediction\"\n",
    "#     )\n",
    "# )\n",
    "# fig.add_trace(\n",
    "#     go.Bar(\n",
    "#         x=np.array(range(len(WPC_smoothed_per_index))), \n",
    "#         y=np.array(WPC_smoothed_per_index),\n",
    "#         name=\"WPC: Smoothed Prediction\"\n",
    "#     )\n",
    "# )\n",
    "# fig.update_layout(\n",
    "#     title='WPCs of each sample. Trained on first 3 minutes, tested on controlled breathing minute @ 0.133hz',\n",
    "#     barmode='group'\n",
    "# )\n",
    "# fig.update_xaxes(tickvals = list(range(len(WPC_per_index))))\n",
    "# fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transfer Learning, pretrain on best data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# np.argsort(WPC_smoothed_per_index)[::-1]\n",
    "argsort_WPC_smoothed_per_index = [24, 19,  9,  4,  0, 23, 22,  2, 20, 25,  3, 26, 17, 13, 16,  8, 12, 10, 21, 28,  1,  5, 27,  7, 18, 14, 11,  6, 15]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_ecg_rolling_train_stack = []\n",
    "y_ip_train_stack = []\n",
    "TRAIN_SAMPLES = 8\n",
    "\n",
    "for train_idx in tqdm(argsort_WPC_smoothed_per_index[:TRAIN_SAMPLES]):#np.argsort(WPC_smoothed_per_index)[::-1][:4]):\n",
    "\n",
    "    print(\"================================================================================\")\n",
    "    print(\"INDEX:\", train_idx)\n",
    "    # ECG input data\n",
    "    \n",
    "    input_ecg_raw = marsh_dataset[train_idx].ECG().transform(transforms=[\n",
    "        Transforms.Crop(start_time_XGB, end_time_XGB),\n",
    "        Transforms.MinMaxScale(),\n",
    "        lambda x: x-x.mean(),\n",
    "        lambda x: detrend(x),\n",
    "    ])\n",
    "    \n",
    "    # Breathing Rate input and target data\n",
    "    input_ip_raw = marsh_dataset[train_idx].IP().transform(transforms=[\n",
    "        Transforms.Crop(start_time_XGB, end_time_XGB),\n",
    "        Transforms.MinMaxScale(),\n",
    "        lambda x: x-x.mean(),\n",
    "        lambda x: detrend(x),\n",
    "    ])\n",
    "    \n",
    "    # rolling window parameters\n",
    "    window_size = input_ecg_raw.sample_rate*4\n",
    "    \n",
    "    input_ecg = input_ecg_raw.transformed_data\n",
    "    input_ip = input_ip_raw.transformed_data\n",
    "    \n",
    "    X_ecg_rolling = get_rolling_windows(input_ecg, window_size=window_size)\n",
    "    y_ip = input_ip.data[window_size-1:]\n",
    "    \n",
    "    X_ecg_rolling_train_stack.append(X_ecg_rolling)\n",
    "    y_ip_train_stack.append(y_ip)\n",
    "\n",
    "X_ecg_rolling_train_stack_np = np.vstack(X_ecg_rolling_train_stack)\n",
    "y_ip_train_stack_np = np.stack(y_ip_train_stack).flatten()\n",
    "\n",
    "# reg = xgb.XGBRegressor(tree_method=\"hist\")\n",
    "# # Fit the model using predictor X and response y.\n",
    "# reg.fit(X_ecg_rolling_train_stack_np, y_ip_train_stack_np)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_ecg_rolling_train_stack_np.shape, y_ip_train_stack_np.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open('asdfasdfX_ecg_rolling_all_stack_np.npy', 'wb') as f:\n",
    "#     np.save(f, asdfasdfafX_ecg_rolling_train_stack_np)\n",
    "# with open('asdfasdfadsy_ip_all_stack_np.npy', 'wb') as f:\n",
    "#     np.save(f, asdasdfay_ip_train_stack_np)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Architecture: A TRANSFORMER-BASED FRAMEWORK FOR MULTIVARIATE TIME SERIES REPRESENTATION LEARNING\n",
    "[Source to Code](https://timeseriesai.github.io/tsai/models.tst.html)\n",
    "\n",
    "[Source to Paper](https://arxiv.org/pdf/2010.02803.pdf)\n",
    "\n",
    "#### Key Points:\n",
    "> At the core of our method lies a transformer encoder, as described in the original transformer work by Vaswani et al. (2017); however, we do not use the decoder part of the architecture\n",
    "\n",
    "\n",
    "> We clarify that regression in the context of this work means predicting a numeric value for a given sequence (time series sample). This numeric\n",
    "value is of a different nature than the numerical data appearing in the time series: for example, given a sequence of simultaneous temperature and humidity measurements of 9 rooms in a house, as well as weather and climate data such as temperature, pressure, humidity, wind speed, visibility and dewpoint, we wish to predict the total energy consumption in kWh of a house for that day\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tsai.all import TST, Learner, MSELossFlat, rmse, ShowGraph, get_splits, \\\n",
    "TSRegression, \\\n",
    "TSStandardize, \\\n",
    "get_ts_dls\n",
    "\n",
    "# Your data preparation here\n",
    "n_variables = 1  # Example number of variables/features in your time series\n",
    "seq_len = 256 #, X_ecg_rolling_train_stack_np.shape[1]   # Length of your time series\n",
    "n_outputs = 1    # Since it's regression, we'll predict a single value\n",
    "iters = 15\n",
    "\n",
    "tsai_X_train = X_ecg_rolling_train_stack_np[:, -seq_len:].reshape(-1, 1, seq_len)\n",
    "tsai_y_train = y_ip_train_stack_np.reshape(-1, 1)\n",
    "# Create TSDataset and TSDataLoaders\n",
    "splits = get_splits(tsai_y_train, valid_size=0.2, stratify=True, random_state=23, shuffle=True)\n",
    "tfms  = [None, TSRegression()]\n",
    "batch_tfms = TSStandardize(by_sample=True)  # Standardize data\n",
    "dls = get_ts_dls(tsai_X_train, tsai_y_train, splits=splits, tfms=tfms, batch_tfms=batch_tfms)\n",
    "\n",
    "# Use a specific architecture directly, e.g., TST for ResNet-like behavior\n",
    "model = TST(dls.vars, dls.c, seq_len=seq_len)\n",
    "learn = Learner(dls, model, loss_func=MSELossFlat(), metrics=rmse, cbs=ShowGraph())\n",
    "\n",
    "# Train the model\n",
    "learn.fit_one_cycle(iters, 3e-4)\n",
    "learn.export(f\"clf_seq-len-{seq_len}-iters-{iters}-samples-{TRAIN_SAMPLES}.pkl\") \n",
    "# Get predictions\n",
    "# preds, targets = learn.get_preds()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Slides Bar Chart"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rmse = [0.150372, 0.150050, 0.150249, 0.116928]\n",
    "training_time = [21, 35, 77, 210]\n",
    "seq_len = [16, 32, 64, 128]\n",
    "\n",
    "fig, ax1 = plt.subplots()\n",
    "\n",
    "ax2 = ax1.twinx()\n",
    "ax1.plot(seq_len, rmse, 'o-', color='blue')\n",
    "ax2.plot(seq_len, training_time, 'o-', color='orange')\n",
    "\n",
    "ax1.set_xlabel('Seq Length')\n",
    "ax1.set_ylabel('Error', color='blue')\n",
    "ax2.set_ylabel('Training Time', color='orange')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Meeting March 25th\n",
    "- [x] Read Chris's paper\n",
    "- [x] Look at context length (window length) and study its effects on the prediction performance\n",
    "\n",
    "#### Meeting April 1st\n",
    "- [ ] Generate prediction signals\n",
    "- [ ] Get xgboost working again.\n",
    "- [ ] Get compute set up for faster experimenting.\n",
    "- [ ] Compare generated signals (WPC between the generated signals and the I.P., estimation statistics for the I.F. scikit/scipy.signal find_peaks)\n",
    "\n",
    "### April 8th\n",
    "- [ ] SCINET tutorial course or documentation. Compute CANADA\n",
    "- [ ] Logistics of sign up done\n",
    "- [ ] Move everything to GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner_128 = load_learner(\"clf_seq-len-128.pkl\")\n",
    "learner_256 = load_learner(\"clf_seq-len-256-iters-6-samples-8.pkl\")\n",
    "\n",
    "X_ecg_rolling_test_stack = []\n",
    "y_ip_test_stack = []\n",
    "\n",
    "for test_idx in tqdm(argsort_WPC_smoothed_per_index[17:17+1]):#np.argsort(WPC_smoothed_per_index)[::-1][:4]):\n",
    "\n",
    "    print(\"================================================================================\")\n",
    "    print(\"INDEX:\", test_idx)\n",
    "    # ECG input data\n",
    "    \n",
    "    input_ecg_raw = marsh_dataset[test_idx].ECG().transform(transforms=[\n",
    "        Transforms.Crop(start_time_XGB, end_time_XGB),\n",
    "        Transforms.MinMaxScale(),\n",
    "        lambda x: x-x.mean(),\n",
    "        lambda x: detrend(x),\n",
    "    ])\n",
    "    \n",
    "    # Breathing Rate input and target data\n",
    "    input_ip_raw = marsh_dataset[test_idx].IP().transform(transforms=[\n",
    "        Transforms.Crop(start_time_XGB, end_time_XGB),\n",
    "        Transforms.MinMaxScale(),\n",
    "        lambda x: x-x.mean(),\n",
    "        lambda x: detrend(x),\n",
    "    ])\n",
    "    \n",
    "    # rolling window parameters\n",
    "    window_size = input_ecg_raw.sample_rate*4\n",
    "    \n",
    "    input_ecg = input_ecg_raw.transformed_data\n",
    "    input_ip = input_ip_raw.transformed_data\n",
    "    \n",
    "    X_ecg_rolling = get_rolling_windows(input_ecg, window_size=window_size)\n",
    "    y_ip = input_ip.data[window_size-1:]\n",
    "    \n",
    "    X_ecg_rolling_test_stack.append(X_ecg_rolling)\n",
    "    y_ip_test_stack.append(y_ip)\n",
    "\n",
    "X_ecg_rolling_test_stack_np = np.vstack(X_ecg_rolling_test_stack)\n",
    "y_ip_test_stack_np = np.stack(y_ip_test_stack).flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seq_len = 256\n",
    "tsai_X_test = X_ecg_rolling_test_stack_np[:, -seq_len:].reshape(-1, 1, seq_len)\n",
    "tsai_y_test = y_ip_test_stack_np.reshape(-1, 1)\n",
    "\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "import torch\n",
    "\n",
    "# Assuming tsai_X_test is a NumPy array, first convert it to a torch Tensor\n",
    "tsai_X_test_tensor = torch.tensor(tsai_X_test).float()  # Convert to float32 tensor\n",
    "\n",
    "# Create a Dataset from the tensor\n",
    "test_dataset = TensorDataset(tsai_X_test_tensor)\n",
    "\n",
    "# Create a DataLoader from this Dataset\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=64)  # You can adjust the batch_size as necessary\n",
    "\n",
    "# Get the predictions using the DataLoader\n",
    "test_preds, _ = learner_256.get_preds(dl=test_dataloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(tsai_X_test[0, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "smoothed_pred = np.convolve(test_preds[:, 0], np.ones((250,))/250, mode='same')[:10000]\n",
    "gt = tsai_y_test[:10000, 0]\n",
    "plt.plot((gt-gt.min())/(gt.max()-gt.min()), '-', label='Ground Truth')\n",
    "plt.plot((smoothed_pred-smoothed_pred.min())/(smoothed_pred.max()-smoothed_pred.min()), '-', label='Prediction')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "smoothed_pred = Signal(format='mat', filepath=None, _type='ECG', data=smoothed_pred)\n",
    "gt = Signal(format='mat', filepath=None, _type='ECG', data=gt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_test, PC_norms_test, PC_test = WPC(\n",
    "    smoothed_pred.transform(transforms=[CWT]).transformed_data,\n",
    "    gt.transform(transforms=[CWT]).transformed_data\n",
    "    # freq=freq_band\n",
    ")\n",
    "print(f\"| {sigma} | {0.0:.4f} | {PC_test.mean():.4f} |\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Single Prediction Process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ECG input data\n",
    "x_test_ecg_raw = marsh_dataset[test_idx].ECG().transform(transforms=[\n",
    "    Transforms.Crop(start_time_XGB_test, start_time_XGB_test+N_test//marsh_dataset[test_idx].ECG().sample_rate),\n",
    "    Transforms.MinMaxScale(),\n",
    "    lambda x: x-x.mean(),\n",
    "    lambda x: detrend(x),\n",
    "]).transformed_data\n",
    "\n",
    "x_test_ecg_rolling = get_rolling_windows(x_test_ecg_raw, window_size=window_size)\n",
    "\n",
    "y_test = marsh_dataset[test_idx].IP().transform(transforms=[\n",
    "    Transforms.Crop(start_time_XGB_test, start_time_XGB_test+N_test//marsh_dataset[test_idx].ECG().sample_rate),\n",
    "    Transforms.MinMaxScale(),\n",
    "    lambda x: x-x.mean(),\n",
    "    lambda x: detrend(x),\n",
    "]).transformed_data[window_size-1:]\n",
    "\n",
    "print(x_test_ecg_rolling.shape, y_test.shape)\n",
    "# Breathing Rate input and target data\n",
    "\n",
    "preds = []\n",
    "smoothed_preds = []\n",
    "smoothing_length = 50\n",
    "mean_filter = np.ones(smoothing_length)\n",
    "\n",
    "for xi in tqdm(x_test_ecg_rolling):\n",
    "    pred = reg.predict(xi[np.newaxis, :])[0]\n",
    "    preds.append(\n",
    "        pred\n",
    "    )\n",
    "    filt_order = min(smoothing_length, len(preds))\n",
    "    smoothed_preds.append(\n",
    "        np.dot(\n",
    "            preds[-filt_order:],\n",
    "            mean_filter[-filt_order:]\n",
    "        )/filt_order\n",
    "    )\n",
    "\n",
    "fig = go.Figure()\n",
    "fig.add_trace(\n",
    "    go.Scatter(\n",
    "        y = preds, name='pred'\n",
    "    )\n",
    ")\n",
    "fig.add_trace(\n",
    "    go.Scatter(\n",
    "        y = smoothed_preds, name='smoothed pred'\n",
    "    )\n",
    ")\n",
    "fig.add_trace(\n",
    "    go.Scatter(\n",
    "        y = y_test, name='ground truth'\n",
    "    )\n",
    ")\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Figure showing Inst. Freq. process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = go.Figure()\n",
    "fig.add_trace(\n",
    "    go.Scatter(\n",
    "        x = np.arange(0, candA.shape[0]/250, 1/250),\n",
    "        y = fantasia_dataset[i].ECG_ENV().transform(transforms=[\n",
    "            crop,\n",
    "            Transforms.MinMaxScale(),\n",
    "            lambda x: x-x.mean()\n",
    "        ]).transformed_data,\n",
    "        name='Envelope and Interplotation'\n",
    "    )\n",
    ")\n",
    "fig.add_trace(\n",
    "    go.Scatter(\n",
    "        x = np.arange(0, candA.shape[0]/250, 1/250),\n",
    "        y = fantasia_dataset[i].ECG().transform(transforms=[\n",
    "            crop,\n",
    "            Transforms.MinMaxScale(),\n",
    "            lambda x: x-x.mean()\n",
    "        ]).transformed_data,\n",
    "        name='ECG'\n",
    "    )\n",
    ")\n",
    "fig.add_trace(\n",
    "    go.Scatter(\n",
    "        x = np.arange(0, candA.shape[0]/250, 1/250),\n",
    "        y = fantasia_dataset[i].IP().transform(transforms=[\n",
    "            crop,\n",
    "            Transforms.MinMaxScale(),\n",
    "            lambda x: x-x.mean()\n",
    "        ]).transformed_data,\n",
    "        name='IP'\n",
    "    )\n",
    ")\n",
    "fig.update_layout(plot_bgcolor='white', width=800, height=500)\n",
    "fig.update_xaxes(showline=True, linewidth=1, linecolor='black')\n",
    "fig.update_yaxes(showline=True, linewidth=1, linecolor='black')\n",
    "fig.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
